package com.itmake.sample.websocket;

import javax.websocket.Session;
import java.io.IOException;
import java.util.Collection;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * socket 工具类
 *
 * @author Amjacks
 * @since 2021/01/10
 */
public class WebSocketMapUtil {

    public static ConcurrentMap<String, Session> webSocketMap = new ConcurrentHashMap<>();

    public static void put(String key, Session session){
        webSocketMap.put(key, session);
    }

    public static Session get(String key){
        return webSocketMap.get(key);
    }

    public static void remove(String key){
        webSocketMap.remove(key);
    }

    public static Collection<Session> getValues(){
        return webSocketMap.values();
    }

    /**
     * 发送消息
     * @param session 客户端 session
     * @param message 消息体
     */
    public static void sendMessage(Session session, String message) {
        try {
            session.getBasicRemote().sendText(message);
        } catch (IOException ioException) {
            throw new RuntimeException("向" + session.getId() + "发送预警失败，异常为" + ioException.getMessage());
        }
    }

    /**
     * 向所有客户端发送消息
     * @param message 消息
     */
    public static void sendMessageToAll(String message) {
        Collection<Session> sessions = WebSocketMapUtil.getValues();
        for (Session session : sessions) {
            WebSocketMapUtil.sendMessage(session, message);
        }
    }
}
